import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import cs224n.util.Counter;
import cs224n.util.PriorityQueue;

/**
 * Represents binomial chi-squared classifier
 * 
 * Use a counter for the chiSquareScores for easy sorting based on score
 * @author alecmgo@gmail.com
 *
 */
public class BinomialChi2Classifier extends BinomialClassifier {
  Map<Integer, Counter<String>> chiSquareScores = new HashMap<Integer, Counter<String>>();   
  Set<String> topWords;
  
  final int NUM_TOP_WORDS = 300;
  //final int NUM_TOP_WORDS = 50000;
  final int NUM_TOP_WORDS_DISPLAY = 20;
  
  public BinomialChi2Classifier(MessageIterator mi) {
    super(mi);
  }
  
  void calculateConditionalProbabilities() {
    topWords = getBestWords();
    calculateConditionalProbabilitiesWithVocabulary(topWords);
  }
  
  double getScore(MessageFeatures message, int newsgroup) {
    return getScoreWithVocabulary(message, newsgroup, topWords);
  }
  
  /**
   * Returns a set of words that score in the top NUM_TOP_WORDS of at least one newsgroup
   * @return
   */
  public Set<String> getBestWords() {
    Set<String> topWords = new HashSet<String>();
    for(int newsgroup = 0; newsgroup < newsgroupCount.size(); newsgroup++) {
      Counter<String> scores = new Counter<String>();
      for(String term : termCount.keySet()) {
        scores.setCount(term, calculateChiSquared(term, newsgroup));
      }
      PriorityQueue<String> queue = scores.asPriorityQueue();
      //System.err.println("\nTop words for newsgroup " + newsgroup);
      for(int i = 0; i < NUM_TOP_WORDS; i++) {
        String term = queue.next();
        if(i < NUM_TOP_WORDS_DISPLAY) {
          //System.e.println(term + "\t" + queue.getPriority());
          System.out.print(term + "\t");
        }
        topWords.add(term);
      }
      System.out.println("");
    }
    return topWords;
  }
  
  /**
   * Returns the chi squared calculation per formula 13.19 on page 256
   * @param term
   * @param newsgroup
   * @return
   */
  private double calculateChiSquared(String term, Integer newsgroup) {
    int n11 = (int) termCount.getCount(term, newsgroup);
    int n10 = (int) (termCount.getCounter(term).totalCount() - termCount.getCount(term, newsgroup));
    int n01 = (int) (newsgroupCount.getCount(newsgroup) - termCount.getCount(term, newsgroup));
    int n00 = (int) (newsgroupCount.totalCount() - n11);
    
    double numerator = (n11 + n10 + n01 + n00) * (n11 * n00 - n10 * n01)^2;
    double denominator = (n11 + n01) * (n11 + n10) * (n10 + n00) * (n01 + n00);
    return numerator / denominator;
  }
}